#
# Ephi - simulation of magnetic fields and particles
# Copyright (C) 2007, 2008 Indrek Mandre <indrek(at)mare.ee>
# For more information please see http://www.mare.ee/indrek/
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#

#
# Use sympy to find the analytic forms of derivatives for cirucal loop of current and
# ring of charge. These derivatives can then be used for bicubic interpolation.
#
# As sympy is quite crude and beta we have to use hacks here to get what we want.
# One problem remains though - fields close to the axis of the loop.
# One has to take the limit of the equation with r->0 to get some sanity
# out of it, but sympy can not manage that at the moment. Maybe in couple
# of years or use a different and better CAS.
#

from sympy import *
import sympy
import re

a, r, R = symbols('arR')
p = r**2+R**2+a**2-2*r*R
q = r**2+R**2+a**2+2*r*R
l = r**2+R**2+a**2
s = R**2-r**2-a**2
u = -2*r-2*R
mu = 4*r*R/q
k = sqrt(mu)
K=Function('K')
L=Function('L')
MC = Symbol('MAGNETIC_CONSTANT')
EMOD = Symbol('EMOD')
PI = Symbol('PREC_PI')
rmod = Symbol('rmod')
Br = MC*((2*a)/(sqrt(q)*r))*(L(k)*(l/p) - K(k))
Ba = MC*(2/sqrt(q))*(L(k)*(s/p)+K(k))
Er = EMOD * (2 / (PI * q**Rational(3,2)*(1-mu))) * (1/mu) * (2*R*K(k)*(1-mu)-L(k)*(2*R-mu*(r+R)))
Ea = EMOD * (2 / (PI * q**Rational(3,2)*(1-mu)))*a*L(k)

def elldiff(f, v):
  ret = diff(f, v)
  dK = diff(K(k),a)/(diff(k,a))
  dL = diff(L(k),a)/(diff(k,a))
  #return ret.subs_dict ({dK: L(k) / (k * (1 - k**2)) - K(k) / k, dL: (L(k) - K(k))/k})
  strf = ret.tostr()
  strf = strf.replace (dK.tostr(), '(' + (L(k)/(k*(1-k**2))-K(k)/k).tostr() + ')')
  strf = strf.replace (dL.tostr(), '(' + ((L(k) - K(k))/k).tostr() + ')')
  return sympify(strf)

def mys(f):
  ret = f.subs(mu, Symbol('mu'))
  ret = ret.subs(p, Symbol('p'))
  ret = ret.subs(q, Symbol('q'))
  ret = ret.subs(l, Symbol('l'))
  ret = ret.subs(s, Symbol('s'))
  ret = ret.subs(u, Symbol('u'))
  return ret

def ltex(f):
  str = latex(f)
  str = str.replace ('\operatorname{K}', 'K')
  str = str.replace ('\operatorname{L}', 'E')
  return str[1:-1]

def splitup(out, eq):
  if hasattr(eq, '_args'):
    for a in eq._args:
      if hasattr(a, '_args') and len(a._args) != 0:
        out.append(a)
        splitup(out, a)

def eqreplace(eq, replace, var):
  if hasattr(eq, '_args') and len(eq._args) != 0:
    newargs = []
    for i in eq._args:
      if i == replace:
          newargs.append(var)
      else:
          newargs.append(eqreplace (i, replace, var))
    t = type(eq)
    return t(*tuple(newargs))
  else:
    return eq

def countup(list):
  ret = {}
  for i in list:
    if not ret.has_key(i):
      ret[i] = 0
    ret[i] = ret[i] + 1
  return ret

def longest(hash):
  ret = None
  for (k, v) in hash.items():
    l = len(k.tostr())
    if len(k._args) != 0 and (ret is None or len(ret.tostr()) < l):
      ret = k
  return ret

#print ltex(elldiff(Ba, a));
#exit(0);

# magnetic field
if 0:
  eqs = [
    ('Ba', Ba),
    ('dBada', elldiff(Ba, a)),
    ('dBadr', rmod*elldiff(Ba, r)),
    ('d2Badadr', rmod*elldiff(elldiff(Ba, a), r)),
    ('Br', rmod*Br),
    ('dBrda', rmod*elldiff(Br, a)),
    ('dBrdr', elldiff(Br, r)),
    ('d2Brdadr', elldiff(elldiff(Br, a), r))
  ]

# electric field
if 0:
  eqs = [
    ('Ea', Ea),
    ('dEada', elldiff(Ea, a)),
    ('dEadr', elldiff(Ea, r)),
    ('d2Eadadr', elldiff(elldiff(Ea, a), r)),
    ('Er', Er),
    ('dErda', elldiff(Er, a)),
    ('dErdr', elldiff(Er, r)),
    ('d2Erdadr', elldiff(elldiff(Er, a), r))
  ]

# magnetic field
# on axis case
# XXX really need limit(Br, r, 0), limit(elldiff(Br, a), r, 0), etc.
# but sympy is not up to it, so we are forced to manually cheat later on
if 0:
  Ba = (MC*2*PI*R*R)/((R*R+a*a)**(Rational(3,2)))
  Br = sympify(0)
  eqs = [
    ('Br', Br),
    ('dBrda', elldiff(Br, a)),
    ('dBrdr', elldiff(Br, r)),
    ('d2Brdadr', elldiff(elldiff(Br, a), r)),
    ('Ba', Ba),
    ('dBada', elldiff(Ba,a)),
    ('dBadr', elldiff(Ba,r)),
    ('d2Badadr', elldiff(elldiff(Ba,a),r))
  ]

# electric field
# on axis case
# XXX really need limit(Er, r, 0), limit(elldiff(Er, a), r, 0), etc.
# but sympy is not up to it, so we are forced to manually cheat later on
if 1:
  Ea = EMOD * a / (R*R+a*a) ** Rational(3,2);
  Er = sympify(0)
  eqs = [
    ('Er', Er),
    ('dErda', elldiff(Er, a)),
    ('dErdr', elldiff(Er, r)),
    ('d2Erdadr', elldiff(elldiff(Er, a), r)),
    ('Ea', Ea),
    ('dEada', elldiff(Ea,a)),
    ('dEadr', elldiff(Ea,r)),
    ('d2Eadadr', elldiff(elldiff(Ea,a),r))
  ]

CC = 0

def eqreduce(eqpairs):
  parts = []
  for i in eqpairs:
    splitup(parts, i[1])
  count = countup(parts)
  l = longest(count)
  if l is None:
      return False
  global CC
  vn = 'v%d' % (CC,)
  CC = CC + 1
  for i in range(0, len(eqpairs)):
    eqpairs[i] = (eqpairs[i][0], eqreplace (eqpairs[i][1], l, Symbol(vn)))
  eqpairs.insert (0, (vn, l**1))
  return True

if 1:
  while eqreduce(eqs):
    pass

if 1:
  print "from __future__ import division"

  print """#m=0.0369686, fk=1.58562, ek=1.55618
def K(k):
  print 'K(%g)' % (k*k,)
  return 1.58562371876423013894452651584288

def L(k):
  print 'L(%g)' % (k*k,)
  return 1.55617658872176023443500980647514

MAGNETIC_CONSTANT = 1e-7
EMOD = 8987551787.368175506591796875
PREC_PI = 3.141592653589793238462643383279502884197169399375105820974944592304
a = 3.1
r = 0.1
R = 1.0
"""
  print 'print "dEada=" + str(' + elldiff(Ea, a).tostr() + ')'

  for name, eq in eqs:
      print name + ' = ' + eq.tostr()
      print 'print "' + name + '=%.8e" % (' + name + ',)'

  print '"""'

for name, eq in eqs:
  line = eq.tostr().strip()
  if 1:
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(1/2\)", r"prec_t_sqrt(\1)", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(-1/2\)", r"(1.0/prec_t_sqrt(\1))", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*2", r"(\1*\1)", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(-2\)", r"(1.0/(\1*\1))", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(-3\)", r"(1.0/(\1*\1*\1))", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(-3/2\)", r"(1.0/(\1*prec_t_sqrt(\1)))", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(-5/2\)", r"(1.0/(\1*\1*prec_t_sqrt(\1)))", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(-7/2\)", r"(1.0/(\1*\1*\1*prec_t_sqrt(\1)))", line)
    line = re.sub (r"([a-zA-Z][0-9a-zA-Z]*)\*\*\(-9/2\)", r"(1.0/(\1*\1*\1*\1*prec_t_sqrt(\1)))", line)
    line = re.sub (r"\(([0-9]+)/([0-9]+)\)", r"(\1.0/\2.0)", line)
    line = re.sub (r"^-([0-9]+)/([0-9]+)\*", r"-(\1.0/\2.0)*", line)
    line = re.sub (r"K\(([a-zA-Z][0-9a-zA-Z]*)\)", r"ellipse1(\1*\1)", line)
    line = re.sub (r"L\(([a-zA-Z][0-9a-zA-Z]*)\)", r"ellipse2(\1*\1)", line)
  print 'prec_t ' + name + ' = ' + line + ';'
  #print 'printf ("' + name + '=%.8e\\n", ' + name + ');'

if 1:
  print '"""'

